package jtg.generator;

import com.microsoft.z3.*;
import com.microsoft.z3.Context;
import com.microsoft.z3.Expr;
import jtg.graphics.SootCFG;
import jtg.parser.Z3Parser;
import jtg.state.PathState;
import jtg.visualizer.Visualizer;
import soot.*;
import soot.jimple.*;
import soot.jimple.internal.JAssignStmt;
import soot.jimple.internal.JInstanceFieldRef;
import soot.toolkits.graph.UnitGraph;

import java.io.File;
import java.util.*;

public class AllPathGenerator extends Generator {
    private final List<Unit> unitPath = new ArrayList<>();
    private List<BoolExpr> pathConstraints;
    private Z3Parser parser;
    private Solver solver;
    private final int maxDepth = 100;

    Set<Unit> visitedNode = new HashSet<>();

    public AllPathGenerator(String className, String methodName) {
        super(className, methodName);
        String defaultClsPath = System.getProperty("user.dir") + File.separator + "target" + File.separator + "classes";
        new PrimeCoverageGenerator(defaultClsPath, className, methodName);
    }

    public AllPathGenerator(String classPath, String className, String methodName) {
        super(classPath, className, methodName);
    }

    @Override
    void drawCFG() {
        drawCFG(clsName + "." + mtdName, false);
    }

    public void drawCFG(String graphName, boolean indexLabel) {
        Visualizer.printCFGDot(graphName, ug, indexLabel);
    }

    private List<Local> getJVars() {
        //Jimple自身增加的Locals，不是被测代码真正的变量
        ArrayList<Local> jimpleVars = new ArrayList<Local>();
        for (Local l : body.getLocals()) {
            if (l.toString().startsWith("$")) jimpleVars.add(l);
        }
        return jimpleVars;
    }

    private List<Local> getParameters() {
        return new ArrayList<>(body.getParameterLocals());
    }

    public List<String> generate() {

        ArrayList<String> testSet = null;

        System.out.println("============================================================================");
        System.out.println("Generating test case inputs for method: " + clsName + "." + mtdName + "()");
        System.out.println("============================================================================");
        try {
            testSet = new ArrayList<String>();
            Context context = new Context();

            if(ug.getHeads().size()>1) throw new Exception("The number of heads is more than 1");
            Unit h = ug.getHeads().get(0);

            Map<String, Expr> paras = new HashMap<>();
            Map<String, Integer> latestAssign = new HashMap<>();
            parser = new Z3Parser(context, paras, latestAssign);
            pathConstraints = new ArrayList<>();
            solver = context.mkSolver();
            parser.solver = solver;
            for (Local local : body.getParameterLocals()) {
                latestAssign.put(local.getName(), 0);
                paras.put(local.getName() + "_0", parser.fromLocal(local, pathConstraints));
            }
            pathConstraints.forEach(solver::add);

            // path = ug.getExtendedBasicBlockPathBetween(h, t);
            calculateWithDFS(context, h, testSet);

        } catch (Exception e) {
            System.err.println("Error in generating test cases: ");
            System.err.println(e.toString());
            e.printStackTrace();
        }
        assert testSet != null;
        if (!testSet.isEmpty()) {
            System.out.println();
            System.out.println("The generated test case inputs:");
            int count = 1;
            for (String tc : testSet) {
                System.out.println("( " + count + " ) " + tc);
                count++;
            }
        }
        return testSet;
    }

    public void calculateWithDFS(Context context, Unit stmt, List<String> testSet) throws Exception {

        System.out.println("Current stmt: " + stmt);
        System.out.println("Current pathConstraints: " + pathConstraints);
//        System.out.println("Solver: " + Arrays.toString(solver.getAssertions()));

        if (solver.check() == Status.UNSATISFIABLE) {
            System.out.println("Unreachable path");
            return;
        } else {

        }

        if (unitPath.size() > maxDepth || stmt instanceof ReturnStmt || stmt instanceof ReturnVoidStmt) {
            System.out.println("Reach the end of the path");
            String tc = getTestSetByModel(context, solver.getModel());
            if (tc != null) {
                testSet.add(tc);
            } else {
                throw new Exception("Path constraints not synchronized with the solver.");
            }
            return;
        }

        List<BoolExpr> previousConstraints = new ArrayList<>(pathConstraints);
        Z3Parser previousParser = new Z3Parser(parser);

        try {
            unitPath.add(stmt);
            solver.push();
            List<BoolExpr> newConstraints = new ArrayList<>();

            if (stmt instanceof AssignStmt) {
                AssignStmt assignStmt = (AssignStmt) stmt;
                Value leftOp = assignStmt.getLeftOp();
                Value rightOp = assignStmt.getRightOp();
                Expr<?> rightExpr = parser.fromSootExpr(rightOp);
                if (rightExpr == null) {
                    // leftOp 设置为对应类型任意值
                    if (leftOp instanceof Local) {
                        parser.assignLocal((Local) leftOp, newConstraints);
                    } else {
                        throw new Exception("Unknown type of leftOp: " + leftOp);
                    }
                    System.err.println("Ignore the statement: " + stmt + " assign arbitrary value to " + leftOp);
                } else {
                    if (leftOp instanceof Local) {
                        Expr<?> newLeft = parser.assignLocal((Local) leftOp, newConstraints);
                        newConstraints.add(parser.mkCastEq(newLeft, rightExpr));
                    } else if (leftOp instanceof ArrayRef) {
                        ArrayRef arrayRef = (ArrayRef) leftOp;
                        Expr array = parser.fromSootExpr(arrayRef.getBase());
                        Expr<?> index = parser.fromSootExpr(arrayRef.getIndex());
                        Expr<?> newArray = context.mkStore(array, index, rightExpr);
                        Expr<?> newLeft = parser.assignLocal((Local) arrayRef.getBase(), newConstraints);
                        newConstraints.add(parser.mkCastEq(newLeft, newArray));
                    } else if (leftOp instanceof InstanceFieldRef) {
                        InstanceFieldRef instanceFieldRef = (InstanceFieldRef) leftOp;
                        Expr<?> newLeft = parser.assignInstanceField(solver, instanceFieldRef, newConstraints);
                        newConstraints.add(parser.mkCastEq(newLeft, rightExpr));
                    } else {
                        throw new Exception("Unknown type of leftOp: " + leftOp.getClass().getName());
                    }
                }
                pathConstraints.addAll(newConstraints);
                newConstraints.forEach(solver::add);
                calculateWithDFS(context, ug.getSuccsOf(stmt).get(0), testSet);
                return;
            }

            if (stmt instanceof IfStmt) {
                IfStmt ifStmt = (IfStmt) stmt;
                Value condition = ifStmt.getCondition();
                Expr<?> conditionExpr = parser.fromSootExpr(condition);
                Unit target1 = ug.getSuccsOf(stmt).get(0); // false
                Unit target2 = ug.getSuccsOf(stmt).get(1); // true
                if (conditionExpr instanceof BoolExpr) {
                    BoolExpr trueCondition = (BoolExpr) conditionExpr;
                    BoolExpr falseCondition = context.mkNot(trueCondition);

                    solver.push();
                    solver.add(trueCondition);
                    pathConstraints.add(trueCondition);
                    calculateWithDFS(context, target2, testSet);
                    pathConstraints.remove(pathConstraints.size() - 1);
                    solver.pop();

                    solver.push();
                    solver.add(falseCondition);
                    pathConstraints.add(falseCondition);
                    calculateWithDFS(context, target1, testSet);
                    pathConstraints.remove(pathConstraints.size() - 1);
                    solver.pop();
                } else {
                    throw new Exception("Unknown type of condition: " + condition);
                }
                return;
            }

            Unit nextStmt = ug.getSuccsOf(stmt).get(0);
            calculateWithDFS(context, nextStmt, testSet);

        } finally {
            unitPath.remove(unitPath.size() - 1);
            pathConstraints = previousConstraints;
            parser = previousParser;
            solver.pop();
        }
    }
}
